#pragma OPENCL EXTENSION cl_amd_printf : enable
#define MAX_NEIGH       128    //by mako
#define N_UNROLL        4     //unroll, future by mako
#define SMALL           1.E-7f //avoid dividing by 0.

//density grid: 
//n_grid[NGX][NGY][NGZ]: n_atoms of a grid
//grid[NGX][NGY][NGZ][MAX_PER_GRID]: atom-ids of a grid
#define MAX_PER_GRID    200 // bymako
#define NGX  16  //by mako
#define NGY  16  //by mako
#define NGZ  16  //by mako
#define TOT_NGRID  (NGX*NGY*NGZ)
#define n_grid_index(i,j,k)   (i*NGY*NGZ + j*NGZ + k)
#define grid_index(i,j,k,l)    ((i*NGY*NGZ + j*NGZ + k)*MAX_PER_GRID+l)

#define SWAP (x,y,tmp)  {tmp=x;x=y;y=tmp;}

//add_f_e_slow: >O(N^2),all,simple for test 
//add_f_e_all: O(N),all

//----------------------------------------------------
__kernel void init_grid(  //init grid to -1 and n_grid to 0
    __global int *grid,   //ids, NULL if -1
    __global int *n_grid  //
    ) 
{
    int iglb = get_global_id(0); 
    int sglb = get_global_size(0);

    int i = iglb;
    while(i<TOT_NGRID){
        n_grid[i] = 0;
        for(int j=0;j<MAX_PER_GRID;++j){
            grid[i*MAX_PER_GRID+j] = -1;
        }
        i += sglb;
    }
}
//---

__kernel void create_grid(  //calculate grid
    const int n_atoms,
    const float4 box_lo,  //
    const float4 box_hi,  //
    __global float4 *r_t,  //[n_atoms] position, type
    __global float4 *rr_g, //[n_atoms]: s0-2:reduced unit, s3:grid-index
    __global int *grid,   //[NGX][NGY][NGZ][MAX_PER_GRID] ids, NULL if -1
    __global int *n_grid, //[NGX][NGY][NGZ] num per grid
    __global int *overflow //output status flag, true(1) if overflow, you may need larger MAX_NEIGH, please init it as false(0) before calling
    ) 
{
    int iglb = get_global_id(0); 
    int sglb = get_global_size(0);
    int iloc = get_local_id(0); 
    int nloc = get_local_size(0);
    float4 box_l = box_hi - box_lo;
    box_l.w = 1.f; //avoid dividing by 0

    int i = iglb;
    while(i<n_atoms){
        float4 ri = r_t[i]; //prefetch
        int igx,igy,igz,idx_ng,idx_g;
        float4 rr = fmod((ri-box_lo)/box_l,1.f); 
        rr = (rr<0)? 1.f-SMALL+rr : rr; //SMALL: avoid single precision overflow (eg. -0.0)
        //if(i<4) printf("ig:%d,rr:%f,%f,%f\n",i,rr.x,rr.y,rr.z);
        //igx = rr.x<0 ? (1.f+rr.x-SMALL)*NGX : rr.x*NGX; //SMALL: avoid single precision overflow
        //igy = rr.y<0 ? (1.f+rr.y-SMALL)*NGY : rr.y*NGY;
        //igz = rr.z<0 ? (1.f+rr.z-SMALL)*NGZ : rr.z*NGZ;
        igx = rr.x*NGX;
        igy = rr.y*NGY;
        igz = rr.z*NGZ;
        idx_ng = n_grid_index(igx,igy,igz);
        rr.w = idx_ng;
        rr_g[i] = rr;
        //if(i<4) printf("ig:%d,igx-y-z:%d,%d,%d,idx:%d\n",i,igx,igy,igz,idx_ng);
        //if(i<4) printf("ig:%d,x-y-z:%f,%f,%f\n",i,ri.x,ri.y,ri.z);
        __global volatile int *p = n_grid+idx_ng; 
        int ni = atomic_inc(p); //return previous value
        //if(i<4) printf("ig:%d,ni:%d\n",i,ni);
        //if(i<4) printf("ig:%d,n_grid[%d]:%d\n",i,idx,n_grid[idx]);
        if(ni<MAX_PER_GRID){
            idx_g = grid_index(igx,igy,igz,ni);
            grid[idx_g] = i;
        }
        else{
            *overflow = 1;
            atomic_dec(p); //return previous value
        }
        ///
        i += sglb;
    }
}
//--
__kernel void neigh_to_int4(
    const int n_atoms,
    __global int *neigh,   //[n_atoms][MAX_NEIGH] neighbor list, <0 means NULL
    __global int *n_neigh, //[n_atoms] real number of neighbor
    __global int4 *neigh4,   //[n_atoms][MAX_NEIGH] neighbor list, <0 means NULL
    __global int *n_neigh4 //[n_atoms] real number of neighbor
    )
{
    int iglb = get_global_id(0); 
    int sglb = get_global_size(0);

    int i = iglb;
    while(i<n_atoms){
        int ni = n_neigh[i];
        int ni4 = ceil(ni/4.f);
        n_neigh4[i] = ni4;
        int4 xj4;
        for(int j=0;j<ni4;j++){
            xj4 = (int4)(-1,-1,-1,-1);
            int nidx = i*MAX_NEIGH + j;

            int idx = j*4;
            xj4.x = neigh[idx];

            idx ++;
            if (idx<ni) xj4.y = neigh[idx];
            
            idx ++;
            if (idx<ni) xj4.z = neigh[idx];

            idx ++; 
            if (idx<ni) xj4.w = neigh[idx];

            neigh4[nidx] = xj4;
        }
        i += sglb;
    }
}
//---
__kernel void create_neigh( 
    const int n_atoms,
    const float r2_cut,  //cutoff of r^2 
    const float4 box_l,  //boxsize, s3:should be no-zero to avoid dividing by 0
    const float4 box_p,  //box periodic, 1:periodic, 0:nonperiodic
    __global float4 *r_t,  //[n_atoms] position, type
    __global float4 *rr_g, //[n_atoms]: s0-2:reduced unit, s3:grid-index
    __global int *grid,   //[NGX][NGY][NGZ][MAX_PER_GRID] ids, NULL if -1
    __global int *n_grid, //[NGX][NGY][NGZ] num per grid
    __global int *neigh,   //[n_atoms][MAX_NEIGH] neighbor list, <0 means NULL
    __global int *n_neigh, //[n_atoms] real number of neighbor
    __global int *overflow //output status flag, true(1) if overflow, you may need larger MAX_NEIGH, please init it as false(0) before calling
    ) 
{
    const float sx = 1.f/NGX; //reduced step in x
    const float sy = 1.f/NGY; //reduced step in x
    const float sz = 1.f/NGZ; //reduced step in x
    int iglb = get_global_id(0); 
    int sglb = get_global_size(0);
    int iloc = get_local_id(0); 
    int nloc = get_local_size(0);

    int i = iglb;
    while(i<n_atoms){
        int ni = 0; //n_neighbor of i
        float4 ri = r_t[i]; //prefetch
        float4 rr = rr_g[i];
        for(int j=0;j<MAX_NEIGH;++j){ //init all to -1
            int idx = j*n_atoms + i;
            neigh[idx] = -1;
        }
        //use mako to unroll:
        for(int ii=-1;ii<2;++ii){
            float fgx = rr.x + ii*sx;
            fgx = fmod(fgx,1.f);
            fgx = (fgx<0)? 1.f-SMALL+fgx : fgx; //SMALL: avoid single precision overflow (eg. -0.0)
            int jgx = fgx*NGX;
            for(int jj=-1;jj<2;++jj){
                float fgy = rr.y + jj*sy;
                fgy = fmod(fgy,1.f);
                fgy = (fgy<0)? 1.f-SMALL+fgy : fgy; //SMALL: avoid single precision overflow (eg. -0.0)
                int jgy = fgy*NGY;
                for(int kk=-1;kk<2;++kk){
                    float fgz = rr.z + kk*sz;
                    fgz = fmod(fgz,1.f);
                    fgz = (fgz<0)? 1.f-SMALL+fgz : fgz; //SMALL: avoid single precision overflow (eg. -0.0)
                    int jgz = fgz*NGZ;
                    //
                    int idx_n_grid = n_grid_index(jgx,jgy,jgz);
                    for(int kj=0;kj<n_grid[idx_n_grid];++kj){
                        int idx_grid = grid_index(jgx,jgy,jgz,kj);
                        int j = grid[idx_grid];
                        float4 rj = r_t[j];
                        float4 d = ri - rj;
                        d = d - round(d/box_l)*box_l*box_p; //periodic or nonperiodic box
                        float r2 = d.x*d.x + d.y*d.y + d.z*d.z;
                        if (r2 < r2_cut && i != j){
                            int idx = ni*n_atoms + i;
                            neigh[idx] = j;
                            ++ni;
                            if(ni>=MAX_NEIGH){
                                *overflow = true;
                                //break;
                                goto over;
                            }
                        }
                    }
                }
            }
        }
over:
        n_neigh[i] = ni;
        i += sglb;
    }
}

//----------------------------------------------------
__kernel void sort_inside_neigh( //slow O^2 sort, but for small size is ok
    const int n_atoms,
    __global int *neigh,   //[n_atoms][MAX_NEIGH] neighbor list, <0 means NULL
    __global int *n_neigh  //[n_atoms] real number of neighbor
    ) 
{
    int iglb = get_global_id(0); 
    int sglb = get_global_size(0);

    int i = iglb;
    while(i<n_atoms){
        int ni = n_neigh[i];
        for(int ii=ni-2;ii>=0;ii--){
            for(int jj=0;jj<ii;jj++){
                int idx = jj*n_atoms + i;
                int tmp;
                int djj =  abs(neigh[idx]-i);
                int djja = abs(neigh[idx+n_atoms]-i);
                //if (neigh[idx]>neigh[idx+1]){
                if (djj>djja){
                    tmp = neigh[idx];
                    neigh[idx] = neigh[idx+n_atoms];
                    neigh[idx+n_atoms] = tmp;
                }
            }
        }
        i += sglb;
    }
}
//

//----------------------------------------------------
__kernel void create_neigh_slow( 
    const int n_atoms,
    const float r2_cut,  //cutoff of r^2 
    const float4 box_l,  //boxsize, s3:should be no-zero to avoid dividing by 0
    const float4 box_p,  //box periodic, 1:periodic, 0:nonperiodic
    __global float4 *r_t,  //[n_atoms] position, type
    __global int *neigh,   //[n_atoms][MAX_NEIGH] neighbor list, <0 means NULL
    __global int *n_neigh, //[n_atoms] real number of neighbor
    __global int *overflow //output status flag, true(1) if overflow, you may need larger MAX_NEIGH, please init it as false(0) before calling
    ) 
{
    int iglb = get_global_id(0); 
    int sglb = get_global_size(0);
    int iloc = get_local_id(0); 
    int nloc = get_local_size(0);

    int i = iglb;
    while(i<n_atoms){
        int ni = 0; //n_neighbor of i
        float4 xi = r_t[i]; //prefetch
        for(int j=0;j<MAX_NEIGH;++j){ //init all to -1
            int idx = j*n_atoms + i;
            neigh[idx] = -1;
        }
        int j = 0;
        while(j<n_atoms){
          //for(int kk=0;kk<N_UNROLL;kk++){ //let compiler do unrolling 
            float4 xj = r_t[j];
            float4 d = xj - xi;
            d = d - round(d/box_l)*box_l*box_p; //periodic or nonperiodic box
            float r2 = d.x*d.x + d.y*d.y + d.z*d.z;
            if (r2 < r2_cut && i != j){
                int idx = ni*n_atoms + i;
                neigh[idx] = j;
                ++ni;
                if(ni>=MAX_NEIGH){
                    *overflow = true;
                    //break;
                    goto over;
                }
            }
            ++j;
          //} //end-unroll
        } //end-while
        n_neigh[i] = ni;
        //if(i<4) printf("ig:%d,n-neigh:%d\n",i,ni);
        ///
over:
        i += sglb;
    }
}
